import operator
from math import log


def calcShannonEnt(dataSet):
    """
    计算数据集的熵
    :param dataSet:
    :return:
    """
    numEntries = len(dataSet)
    labelCounts = {}

    for featVec in dataSet:
        currentLabel = featVec[-1]
        labelCounts[currentLabel] = labelCounts.get(currentLabel, 0) + 1

    # 经验熵
    shannonEnt = 0
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntries
        shannonEnt -= prob * log(prob, 2)
    return shannonEnt


def createDataSet():
    """
    创建数据集
    """
    dataSet = [[0, 0, 0, 0, 'no'],  # 数据集
               [0, 0, 0, 1, 'no'],
               [0, 1, 0, 1, 'yes'],
               [0, 1, 1, 0, 'yes'],
               [0, 0, 0, 0, 'no'],
               [1, 0, 0, 0, 'no'],
               [1, 0, 0, 1, 'no'],
               [1, 1, 1, 1, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [2, 0, 1, 2, 'yes'],
               [2, 0, 1, 1, 'yes'],
               [2, 1, 0, 1, 'yes'],
               [2, 1, 0, 2, 'yes'],
               [2, 0, 0, 0, 'no']]
    labels = ['年龄', '有工作', '有自己的房子', '信贷情况']  # 特征标签
    return dataSet, labels  # 返回数据集和分类属性


def splitDataSet(dataSet, index, value):
    """
    分隔数据集
    :param dataSet:
    :param index:
    :param value:
    :return:
    """
    retDataSet = []
    for featVec in dataSet:
        if featVec[index] == value:
            reducedFeatVec = featVec[:index]
            reducedFeatVec.extend(featVec[index + 1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet


def chooseBestFeatureToSplit(dataSet):
    """
    选择最优分隔特征
    :param dataSet:
    :return:
    """
    # 特征数量
    numberFeatures = len(dataSet[0]) - 1
    # 计算数据集熵
    baseEntropy = calcShannonEnt(dataSet)
    # 信息增益
    bestInfoGain = 0.0
    # 最优划分特征的索引
    bestFeature = -1

    # 遍历所有特征
    for i in range(numberFeatures):
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)

        # 经验条件熵
        newEntropy = 0
        for value in uniqueVals:
            # subDataSet划分后的子集
            subDataSet = splitDataSet(dataSet, i, value)
            # 计算子集的概率
            prob = len(subDataSet) / float(len(dataSet))
            # 根据公式计算经验条件熵
            newEntropy += prob * calcShannonEnt(subDataSet)
        # 信息增益
        infoGain = baseEntropy - newEntropy

        if (infoGain > bestInfoGain):
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature


def majorityCnt(classList):
    """
    少数服从多数
    :param classList:
    :return:
    """
    classCount = {}
    for vote in classList:
        classCount[vote] = classCount.get(vote, 0) + 1
    sortedCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedCount[0][0]


def createTree(dataSet, labels, featLabels):
    """
    :param dataSet: 训练数据集
    :param labels: 分类属性标签
    :param featLabels: 存储选择的最优特征标签
    :return:
    """
    classList = [example[-1] for example in dataSet]

    # 第一个停止条件：若类别相同，则停止继续划分
    if classList.count(classList[0]) == len(classList):
        return classList[0]

    # 第二个停止条件：遍历完所有特征时返回出现次数最多的类标签
    if len(dataSet[0]) == 1:
        return majorityCnt(classList)

    bestFeat = chooseBestFeatureToSplit(dataSet)
    bestFeatLabel = labels[bestFeat]
    featLabels.append(bestFeatLabel)

    # 根据最优特征标签生成树
    myTree = {bestFeatLabel: {}}
    # 删除已使用过的特征标签
    del (labels[bestFeat])
    # 得到训练集中所有最优特征属性值集合
    featValues = [example[bestFeat] for example in dataSet]
    # 去掉重复属性
    uniqueVals = set(featValues)
    # 遍历特征，创建决策树
    for value in uniqueVals:
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), labels, featLabels)
    return myTree


def classify(inputTree, featLabels, testVec):
    """
    分类
    :param inputTree: 构造好的树（字典类型）
    :param featLabels: 分隔节点特征
    :param testVec: 测试数据集向量
    :return:
    """
    # 取第一个字典的键 key
    firstStr = next(iter(inputTree))
    # 下一个字典
    secondDict = inputTree[firstStr]
    featIndex = featLabels.index(firstStr)
    for key in secondDict.keys():
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:
                classLabel = secondDict[key]
    return classLabel


if __name__ == "__main__":
    myDat, labels = createDataSet()
    featLables = []

    myTree = createTree(myDat, labels, featLables)
    # 测试数据
    testVec = [0, 0, 0, 0]
    result = classify(myTree, featLables, testVec)
    if result == 'yes':
        print('放贷')
    if result == 'no':
        print('不放贷')
